RED = 1
BLACK = 0


class red_bleak_tree_node(object):
    def __init__(self, value):
        self.value = value
        self.right = None
        self.left = None
        self.p = None
        self.color = RED


# 左旋  node是要旋转的节点，tree是self自身
def left_rotate(tree, node: red_bleak_tree_node):
    if not node.right:
        return False
    node_right = node.right
    node_right.p = node.p
    if not node.p:
        tree.root = node_right
    elif node.p.left == node:
        node.p.left = node_right
    else:
        node.p.right = node_right
    if node_right.left:
        node_right.left.p = node
    node.right = node_right.left
    node.p = node_right
    node_right.left = node


# 右旋，node同上
def right_rotate(tree, node: red_bleak_tree_node):
    if not node.left:
        return False
    node_left = node.left
    node_left.p = node.p
    if not node.p:
        tree.root = node_left
    elif node.p.right == node:
        node.p.right = node_left
    elif node.p.left == node:
        node.p.left = node_left
    if node_left.right:
        node_left.right.p = node
    node.left = node_left.right

    node.p = node_left
    node_left.right = node

class red_bleak_tree(object):
    def __init__(self):
        self.root = None

    def insert(self, node: red_bleak_tree_node):
        if not self.root :
            self.root = node
        else:
            x = self.root
            while x:
                y = x
                if x.value > node.value:
                    x = x.left
                else:
                    x = x.right
            node.p = y
            if node.value>y.value:
                y.right = node
            else:
                y.left = node
        self.insert_fixup(node)

    def insert_fixup(self, node: red_bleak_tree_node):
        parent: red_bleak_tree_node = node.p
        while parent and parent.color == RED:
            grand: red_bleak_tree_node = parent.p
            if parent == grand.left:
                uncle: red_bleak_tree_node = grand.right
                # 判断情形3
                if uncle and uncle.color == RED:
                    grand.color = RED
                    parent.color = BLACK
                    uncle.color = BLACK
                    node = grand
                    parent = node.p
                    continue

                    # 情形四
                if node == parent.right:
                    left_rotate(self, parent)
                    node, parent = parent, node

                # 情形五
                right_rotate(self, grand)
                grand.color = RED
                parent.color = BLACK
                break
            else:
                uncle: red_bleak_tree_node = grand.left
                # 判断情形3
                if uncle and uncle.color == RED:
                    grand.color = RED
                    parent.color = BLACK
                    uncle.color = BLACK
                    node = grand
                    parent = node.p
                    continue

                    # 情形四
                if node == parent.left:
                    right_rotate(self, parent)
                    node, parent = parent, node

                # 情形五
                left_rotate(self, grand)
                grand.color = RED
                parent.color = BLACK
                break
        self.root.color = BLACK

def rbtree_print(node, key, direction):
    if node:
        if direction == 0:  # tree 是根节点
            print("%2d(B) is root" % node.value)
        else:  # tree 是分支节点
            print("%2d(%s) is %2d's %6s child" % (
                node.value, ("B" if node.color == 1 else "R"), key,
                ("right" if direction == 1 else "left")))
        rbtree_print(node.left, node.value, -1)
        rbtree_print(node.right, node.value, 1)


def main():
    number_list = (7, 4, 1, 8, 5, 2, 9, 6, 3)
    tree = red_bleak_tree()
    for number in number_list:
        node_1 = red_bleak_tree_node(number)
        tree.insert(node_1)
    del node_1
    rbtree_print(tree.root, tree.root.value, 0)


if __name__ == '__main__':
    main()
